import argparse
import os
import shutil
import random
import glob
import copy
import numpy as np
import cv2
import multiprocessing
import math
import pprint

import torch
import yaml
from imutils import paths
from sklearn.model_selection import train_test_split
from tqdm import tqdm

import lcnn
from dataset.constants import TB_DATATYPE, LR_DATATYPE
from demo import get_model_eval
from lcnn.config import C, M
from dataset.train_test_split import get_json_pts
from lcnn.models.line_vectorizer import LineVectorizer
from lcnn.models.multitask_learner import MultitaskHead, MultitaskLearner
from python_developer_tools.cv.datasets.line_imgaug import random_perspective
from python_developer_tools.cv.utils.torch_utils import init_cudnn, init_seeds

# perspective_r: 0.5 #进行下面的增强概率
# degrees: 0.1  # image rotation (+/- deg)
# translate: 0.1  # image translation (+/- fraction)
# scale: 0.5  # image scale (+/- gain) 缩放
# shear: 0.1  # image shear (+/- deg) 旋转
# perspective: 0.00  # image perspective (+/- fraction), range 0-0.001 透视图


# img,label = random_perspective(img,label,
#                          degrees=self.hyp['degrees'],
#                          translate=self.hyp['translate'],
#                          scale=self.hyp['scale'],
#                          shear=self.hyp['shear'],
#                          perspective=self.hyp['perspective'])


if __name__ == '__main__':
    org_dir = r'/home/zengxh/medias/data/ext/creepageDistance/lab_datasets/lr/org'
    save_dir = r'/home/hongle/Temp/lcnn/0804/drawing-lr'
    image_paths = list(paths.list_images(org_dir))
    for img_path in image_paths:
        line_type = TB_DATATYPE
        path = img_path.replace('.jpg', '.json')
        _, img_name = os.path.split(img_path)
        img_type = img_name.split('.')[0][-1]
        if img_type == 'l' or img_type == 'r':
            line_type = LR_DATATYPE

        # 获取真实标注数据点
        true_pts = []
        true_pts = get_json_pts(path)
        true_pts = np.array(true_pts)

        true_pts_reshape = []
        for i in range(len(true_pts)):
            true_pts_reshape.append(true_pts[i].flatten())

            # true_pts = np.array(true_pts).reshape(4, 2)
            # true_pts = true_pts[:, ::-1]

        org_img = cv2.imread(img_path)
        # _ = cv2.line(org_img, (int(true_pts_reshape[0][0]), int(true_pts_reshape[0][1])), (int(true_pts_reshape[0][2]), int(true_pts_reshape[0][3])), (0, 0, 255), thickness=2)
        # _ = cv2.line(org_img, (int(true_pts_reshape[1][0]), int(true_pts_reshape[1][1])),
        #                 (int(true_pts_reshape[1][2]), int(true_pts_reshape[1][3])), (0, 0, 255), thickness=2)
        # org_img_name = img_name.split('.')[0] + '_org' + '.jpg'
        # cv2.imwrite(os.path.join(save_dir, org_img_name), org_img)

        transformed_img, transformed_pts, is_transformed = random_perspective(org_img, true_pts_reshape)
        print("sdfasd")
        # if is_transformed:
        #     _ = cv2.line(transformed_img, (int(transformed_pts[0][0]), int(transformed_pts[0][1])),
        #                  (int(transformed_pts[0][2]), int(transformed_pts[0][3])), (0, 0, 255), thickness=2)
        #     _ = cv2.line(transformed_img, (int(transformed_pts[1][0]), int(transformed_pts[1][1])),
        #                  (int(transformed_pts[1][2]), int(transformed_pts[1][3])), (0, 0, 255), thickness=2)
        #     transformed_img_name = img_name.split('.')[0] + '_transformed' + '.jpg'
        #     cv2.imwrite(os.path.join(save_dir, transformed_img_name), transformed_img)



